package com.hevery.warpper;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;

/**
 * @Author: Hevery
 * @Description: 对传入的ServletResponse进行包装，
 * @Date: Created on 2017/10/8
 * @Modified By:
 */

public class HttpServletResponseCopier extends HttpServletResponseWrapper {

    private ServletOutputStream outputStream;
    private PrintWriter writer;
    private ServletOutputStreamCopier copier;

    public HttpServletResponseCopier(HttpServletResponse response) throws IOException {
        super(response);
    }


    /**获得对应的字节流，getResponse().getOutputStream()可以得到构造中传入的HttpServletResponse的字节流*/
    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        if (writer != null) {
            throw new IllegalStateException("getWriter() has already been called");
        }
        /**核心代码*/
        if (outputStream == null) {
            outputStream = getResponse().getOutputStream();
            /**此处将传入的HttpServletResponse的字节流进行包装，转化为我们自定义的ServletOutputStreamCopier流*/
            copier = new ServletOutputStreamCopier(outputStream);
        }
        return copier;
    }
    /**将对应的字节流转换为字符流*/
    @Override
    public PrintWriter getWriter() throws IOException {
        if (outputStream != null) {
            throw new IllegalStateException("getOutputStream() has been called");
        }
        if (writer == null) {
            copier = new ServletOutputStreamCopier(getResponse().getOutputStream());
            writer = new PrintWriter(new OutputStreamWriter(copier, getResponse().getCharacterEncoding()), true);
        }
        return writer;
    }

    @Override
    public void flushBuffer() throws IOException {
        if (writer != null) {
            writer.flush();
        } else if (outputStream != null) {
            copier.flush();
        }
    }

    /**获得ServletOutputStreamCopier流中的ByteArrayOutputStream(对ServletResponse流的拷贝)*/
    public byte[] getCopy() {
        if (copier != null) {
            return copier.getCopy();
        } else {
            return new byte[0];
        }
    }
}
